/*
 * MIT License
 *
 * Copyright (c) 2023-present, tangli
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in all
 * copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 * SOFTWARE.
 */

package tony.crypto.asymmetric

import java.security.KeyFactory
import java.security.KeyPair
import java.security.KeyPairGenerator
import java.security.PrivateKey
import java.security.PublicKey
import java.security.spec.PKCS8EncodedKeySpec
import java.security.spec.X509EncodedKeySpec
import org.bouncycastle.asn1.ASN1EncodableVector
import org.bouncycastle.asn1.ASN1OctetString
import org.bouncycastle.asn1.ASN1Sequence
import org.bouncycastle.asn1.DERSequence
import org.bouncycastle.crypto.engines.SM2Engine
import org.bouncycastle.crypto.params.ECDomainParameters
import org.bouncycastle.crypto.params.ECPrivateKeyParameters
import org.bouncycastle.crypto.params.ECPublicKeyParameters
import org.bouncycastle.crypto.params.ParametersWithRandom
import org.bouncycastle.crypto.signers.SM2Signer
import org.bouncycastle.jcajce.provider.asymmetric.ec.BCECPrivateKey
import org.bouncycastle.jcajce.provider.asymmetric.ec.BCECPublicKey
import org.bouncycastle.jce.ECNamedCurveTable
import org.bouncycastle.jce.spec.ECNamedCurveGenParameterSpec
import org.bouncycastle.jce.spec.ECNamedCurveParameterSpec
import org.bouncycastle.jce.spec.ECPublicKeySpec
import org.bouncycastle.math.ec.ECPoint
import tony.core.utils.secureRandom
import tony.crypto.bouncyCastleProvider

/**
 * SM2算法实例
 *
 * 提供SM2算法的加密、解密、签名、验证等功能
 * 支持C1C2C3和C1C3C2两种密文格式
 *
 * @param mode 密文格式模式，默认为C1C3C2（国密标准）
 * @author tony
 * @since 2024-12-20
 */
public class Sm2Instance(
    private var mode: SM2Engine.Mode = SM2Engine.Mode.C1C2C3,
) {
    /**
     * SM2算法标识符
     */
    public companion object {
        private const val ALGORITHM = "EC"
        private const val CURVE_NAME = "sm2p256v1"
    }

    private val sM2Engine: SM2Engine by lazy {
        SM2Engine(mode)
    }

    private val signer: SM2Signer by lazy {
        SM2Signer()
    }

    private val keyFactory: KeyFactory by lazy {
        KeyFactory.getInstance(ALGORITHM, bouncyCastleProvider)
    }

    private val keyPairGenerator: KeyPairGenerator by lazy {
        KeyPairGenerator.getInstance(ALGORITHM, bouncyCastleProvider)
    }

    private val ecNamedCurveParameterSpec: ECNamedCurveParameterSpec by lazy {
        ECNamedCurveTable.getParameterSpec(CURVE_NAME)
    }

    /**
     * 生成SM2密钥对
     *
     * @return SM2密钥对
     * @throws Exception 密钥生成失败时抛出异常
     */
    public fun generateKeyPair(): KeyPair {
        keyPairGenerator.initialize(ECNamedCurveGenParameterSpec(CURVE_NAME), secureRandom)
        return keyPairGenerator.generateKeyPair()
    }

    /**
     * SM2加密，输出配置的密文格式
     *
     * @param plaintext 明文数据
     * @param publicKeyBytes 公钥字节数组（X.509 编码或65 字节未压缩椭圆曲线点格式）
     * @return 密文
     * @throws Exception 加密失败时抛出异常
     */
    public fun encrypt(
        plaintext: ByteArray,
        publicKeyBytes: ByteArray,
    ): ByteArray {
        val publicKey = parsePublicKey(publicKeyBytes)
        return encrypt(plaintext, publicKey)
    }

    /**
     * SM2加密 - 使用PublicKey对象
     *
     * @param plaintext 明文数据
     * @param publicKey 公钥对象
     * @return 密文
     * @throws Exception 加密失败时抛出异常
     */
    public fun encrypt(
        plaintext: ByteArray,
        publicKey: PublicKey,
    ): ByteArray {
        val bcPublicKey = publicKey as BCECPublicKey
        val ecPublicKeyParameters =
            ECPublicKeyParameters(
                bcPublicKey.q,
                ECDomainParameters(
                    bcPublicKey.parameters.curve,
                    bcPublicKey.parameters.g,
                    bcPublicKey.parameters.n
                )
            )

        sM2Engine.init(true, ParametersWithRandom(ecPublicKeyParameters, secureRandom))
        return sM2Engine.processBlock(plaintext, 0, plaintext.size)
    }

    /**
     * SM2解密
     *
     * @param ciphertext 密文
     * @param privateKeyBytes 私钥字节数组（PKCS8编码）
     * @return 解密后的明文
     * @throws Exception 解密失败时抛出异常
     */
    public fun decrypt(
        ciphertext: ByteArray,
        privateKeyBytes: ByteArray,
    ): ByteArray {
        val privateKey = parsePrivateKey(privateKeyBytes)
        return decrypt(ciphertext, privateKey)
    }

    /**
     * SM2解密 - 使用PrivateKey对象
     *
     * @param ciphertext 密文
     * @param privateKey 私钥对象
     * @return 解密后的明文
     * @throws Exception 解密失败时抛出异常
     */
    public fun decrypt(
        ciphertext: ByteArray,
        privateKey: PrivateKey,
    ): ByteArray {
        val bcPrivateKey = privateKey as BCECPrivateKey
        val ecPrivateKeyParameters =
            ECPrivateKeyParameters(
                bcPrivateKey.d,
                ECDomainParameters(
                    bcPrivateKey.parameters.curve,
                    bcPrivateKey.parameters.g,
                    bcPrivateKey.parameters.n
                )
            )

        sM2Engine.init(false, ecPrivateKeyParameters)
        return sM2Engine.processBlock(ciphertext, 0, ciphertext.size)
    }

    /**
     * SM2签名，使用SM3withSM2算法
     *
     * @param data 待签名数据
     * @param privateKeyBytes 私钥字节数组（PKCS8编码）
     * @return 签名结果
     * @throws Exception 签名失败时抛出异常
     */
    public fun sign(
        data: ByteArray,
        privateKeyBytes: ByteArray,
    ): ByteArray {
        val privateKey = parsePrivateKey(privateKeyBytes)
        return sign(data, privateKey)
    }

    /**
     * SM2签名 - 使用PrivateKey对象
     *
     * @param data 待签名数据
     * @param privateKey 私钥对象
     * @return 签名结果
     * @throws Exception 签名失败时抛出异常
     */
    public fun sign(
        data: ByteArray,
        privateKey: PrivateKey,
    ): ByteArray {
        val bcPrivateKey = privateKey as BCECPrivateKey
        val ecPrivateKeyParameters =
            ECPrivateKeyParameters(
                bcPrivateKey.d,
                ECDomainParameters(
                    bcPrivateKey.parameters.curve,
                    bcPrivateKey.parameters.g,
                    bcPrivateKey.parameters.n
                )
            )

        signer.init(true, ecPrivateKeyParameters)
        signer.update(data, 0, data.size)
        return signer.generateSignature()
    }

    /**
     * SM2验签，使用SM3withSM2算法
     *
     * @param data 原始数据
     * @param signature 签名数据
     * @param publicKeyBytes 公钥字节数组（X.509 编码或65 字节未压缩椭圆曲线点格式）
     * @return 验签结果，true表示验签成功
     * @throws Exception 验签失败时抛出异常
     */
    public fun verify(
        data: ByteArray,
        signature: ByteArray,
        publicKeyBytes: ByteArray,
    ): Boolean {
        val publicKey = parsePublicKey(publicKeyBytes)
        return verify(data, signature, publicKey)
    }

    /**
     * SM2验签 - 使用PublicKey对象
     *
     * @param data 原始数据
     * @param signature 签名数据
     * @param publicKey 公钥对象
     * @return 验签结果，true表示验签成功
     * @throws Exception 验签失败时抛出异常
     */
    public fun verify(
        data: ByteArray,
        signature: ByteArray,
        publicKey: PublicKey,
    ): Boolean {
        val bcPublicKey = publicKey as BCECPublicKey
        val ecPublicKeyParameters =
            ECPublicKeyParameters(
                bcPublicKey.q,
                ECDomainParameters(
                    bcPublicKey.parameters.curve,
                    bcPublicKey.parameters.g,
                    bcPublicKey.parameters.n
                )
            )

        signer.init(false, ecPublicKeyParameters)
        signer.update(data, 0, data.size)
        return signer.verifySignature(signature)
    }

    /**
     * 将C1C2C3格式转换为C1C3C2格式
     *
     * @param c1c2c3Data C1C2C3格式的数据
     * @return C1C3C2格式的数据
     * @throws IllegalArgumentException 当数据格式无效时抛出异常
     */
    public fun convertC1C2C3ToC1C3C2(c1c2c3Data: ByteArray): ByteArray {
        val sequence = ASN1Sequence.getInstance(c1c2c3Data)
        if (sequence.size() != 3) {
            throw IllegalArgumentException(
                "Invalid SM2 ciphertext format: expected 3 components, got ${sequence.size()}"
            )
        }

        val c1 = sequence.getObjectAt(0) as ASN1OctetString
        val c2 = sequence.getObjectAt(1) as ASN1OctetString
        val c3 = sequence.getObjectAt(2) as ASN1OctetString

        val vector = ASN1EncodableVector()
        vector.add(c1)
        vector.add(c3)
        vector.add(c2)

        return DERSequence(vector).encoded
    }

    /**
     * 将C1C3C2格式转换为C1C2C3格式
     *
     * @param c1c3c2Data C1C3C2格式的数据
     * @return C1C2C3格式的数据
     * @throws IllegalArgumentException 当数据格式无效时抛出异常
     */
    public fun convertC1C3C2ToC1C2C3(c1c3c2Data: ByteArray): ByteArray {
        val sequence = ASN1Sequence.getInstance(c1c3c2Data)
        if (sequence.size() != 3) {
            throw IllegalArgumentException(
                "Invalid SM2 ciphertext format: expected 3 components, got ${sequence.size()}"
            )
        }

        val c1 = sequence.getObjectAt(0) as ASN1OctetString
        val c3 = sequence.getObjectAt(1) as ASN1OctetString
        val c2 = sequence.getObjectAt(2) as ASN1OctetString

        val vector = ASN1EncodableVector()
        vector.add(c1)
        vector.add(c2)
        vector.add(c3)

        return DERSequence(vector).encoded
    }

    /**
     * 解析公钥字节数组为PublicKey对象
     *
     * @param publicKeyBytes 公钥字节数组（X.509 编码或65 字节未压缩椭圆曲线点格式）
     * @return PublicKey对象
     * @throws Exception 解析失败时抛出异常
     */
    public fun parsePublicKey(publicKeyBytes: ByteArray): PublicKey {
        // 尝试解析为X.509 格式
        try {
            val keySpec = X509EncodedKeySpec(publicKeyBytes)
            return keyFactory.generatePublic(keySpec)
        } catch (_: Exception) {
            // 如果X.509 格式失败，尝试解析为原始椭圆曲线点格式
            return parseRawECPublicKey(publicKeyBytes)
        }
    }

    /**
     * 解析原始椭圆曲线点格式的公钥
     *
     * @param keyBytes 原始椭圆曲线点字节数组
     * @return PublicKey对象
     * @throws IllegalArgumentException 当密钥格式不支持时抛出异常
     */
    private fun parseRawECPublicKey(keyBytes: ByteArray): PublicKey {
        // 检查是否为未压缩格式的椭圆曲线点 (04 + 32字节x + 32字节y = 65 字节)
        if (keyBytes.size == 65 && keyBytes[0] == 0x04.toByte()) {
            val point: ECPoint = ecNamedCurveParameterSpec.curve.decodePoint(keyBytes)
            val keySpec = ECPublicKeySpec(point, ecNamedCurveParameterSpec)
            return keyFactory.generatePublic(keySpec)
        }

        throw IllegalArgumentException(
            "不支持的公钥格式，期望X.509编码或65字节的未压缩椭圆曲线点格式，实际长度: ${keyBytes.size}"
        )
    }

    /**
     * 解析私钥字节数组为PrivateKey对象
     *
     * @param privateKeyBytes 私钥字节数组（PKCS8编码）
     * @return PrivateKey对象
     * @throws Exception 解析失败时抛出异常
     */
    public fun parsePrivateKey(privateKeyBytes: ByteArray): PrivateKey {
        val keySpec = PKCS8EncodedKeySpec(privateKeyBytes)
        return keyFactory.generatePrivate(keySpec)
    }
}
